import gradio as gr
import cv2
import numpy as np

def start_box_handler(image, evt: gr.SelectData):
    """处理矩形框开始绘制事件"""
    if image is None or evt is None or not hasattr(evt, '_data') or evt._data is None:
        return image, True, (0, 0)
    orig_x, orig_y = evt.index[0], evt.index[1]
    if orig_x is None or orig_y is None:
        return image, True, (0, 0)
    return image, True, (int(orig_x), int(orig_y))

def update_box_handler(image, is_drawing, start, evt: gr.SelectData):
    """更新矩形框绘制状态"""
    if not is_drawing or image is None or start is None or evt._data is None:
        return image, is_drawing, start
    orig_x, orig_y = evt.index[0], evt.index[1]
    if orig_x is None or orig_y is None:
        return image, is_drawing, start
    
    # Get coordinates as integers
    start_x, start_y = start
    end_x, end_y = orig_x, orig_y
    
    temp = image.copy()
    # Draw rectangle preview
    cv2.rectangle(temp, (int(start_x), int(start_y)), (int(end_x), int(end_y)), (0, 255, 0), 2)
    return temp, is_drawing, start

def end_box_handler(image, is_drawing, start, points_list, boxes_list, evt: gr.SelectData):
    """结束矩形框绘制并添加到列表"""
    if not is_drawing or image is None or start is None or evt._data is None:
        return image, is_drawing, start, boxes_list
    
    orig_x, orig_y = evt.index[0], evt.index[1]
    if orig_x is None or orig_y is None:
        return image, is_drawing, start, boxes_list
    
    # Get coordinates as integers
    start_x, start_y = start
    end_x, end_y = orig_x, orig_y
    
    # Add box coordinates to list
    boxes_list.append((min(int(start_x), int(end_x)), min(int(start_y), int(end_y)), 
                      max(int(start_x), int(end_x)), max(int(start_y), int(end_y))))
    
    temp = image.copy()
    
    # Draw all points
    for pt in points_list:
        cv2.drawMarker(
            temp, 
            (pt[0], pt[1]), 
            (0,0,255), 
            cv2.MARKER_STAR, 
            20, 
            2
        )
        
    # Draw all boxes
    for box in boxes_list:
        x1, y1, x2, y2 = box
        cv2.rectangle(temp, (int(x1), int(y1)), (int(x2), int(y2)), (0, 255, 0), 2)
    
    return temp, False, None, boxes_list

def remove_box_handler(image, boxes_list, evt: gr.SelectData):
    """删除最近的矩形框"""
    if image is None or not boxes_list or evt._data is None:
        return image, boxes_list
    
    # Get click position
    orig_x, orig_y = evt.index[0], evt.index[1]
    if orig_x is None or orig_y is None:
        return image, boxes_list
    
    # Find nearest box
    nearest_idx = min(
        range(len(boxes_list)),
        key=lambda i: min(
            (boxes_list[i][0] - orig_x)**2 + (boxes_list[i][1] - orig_y)**2,
            (boxes_list[i][2] - orig_x)**2 + (boxes_list[i][3] - orig_y)**2
        )
    )
    boxes_list.pop(nearest_idx)
    
    # Redraw image
    temp = image.copy()
    for box in boxes_list:
        x1, y1, x2, y2 = box
        cv2.rectangle(temp, (int(x1), int(y1)), (int(x2), int(y2)), (0, 255, 0), 2)
    
    return temp, boxes_list

def clear_boxes_handler(clean_img):
    return clean_img, []
